package test.lucene.pinyin;

import java.io.IOException;
import org.apache.lucene.analysis.TokenFilter;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.analysis.tokenattributes.OffsetAttribute;

public class YidaPinyinNGramTokenFilter extends TokenFilter
{
  private final int minGram;
  private final int maxGram;
  private final boolean nGramChinese;
  private final boolean nGramNumber;
  private final CharTermAttribute termAtt;
  private final OffsetAttribute offsetAtt;
  private char[] curTermBuffer;
  private int curTermLength;
  private int curGramSize;
  private int tokStart;

  public YidaPinyinNGramTokenFilter(TokenStream input)
  {
    this(input, 2, 20, 
      false);
  }

  public YidaPinyinNGramTokenFilter(TokenStream input, int maxGram) {
    this(input, 2, maxGram, false);
  }

  public YidaPinyinNGramTokenFilter(TokenStream input, int minGram, int maxGram) {
    this(input, minGram, maxGram, false, false);
  }

  public YidaPinyinNGramTokenFilter(TokenStream input, int minGram, int maxGram, boolean nGramChinese) {
    this(input, minGram, maxGram, nGramChinese, false);
  }

  public YidaPinyinNGramTokenFilter(TokenStream input, int minGram, int maxGram, boolean nGramChinese, boolean nGramNumber)
  {
    super(input);

    this.termAtt = ((CharTermAttribute)addAttribute(CharTermAttribute.class));
    this.offsetAtt = ((OffsetAttribute)addAttribute(OffsetAttribute.class));

    if (minGram < 1) {
      throw new IllegalArgumentException(
        "minGram must be greater than zero");
    }
    if (minGram > maxGram) {
      throw new IllegalArgumentException(
        "minGram must not be greater than maxGram");
    }
    this.minGram = minGram;
    this.maxGram = maxGram;
    this.nGramChinese = nGramChinese;
    this.nGramNumber = nGramNumber;
  }

  public static boolean containsChinese(String s) {
    if ((s == null) || ("".equals(s.trim())))
      return false;
    for (int i = 0; i < s.length(); ++i) {
      if (isChinese(s.charAt(i))) {
        return true;
      }
    }
    return false;
  }

  public static boolean isChinese(char a) {
    int v = a;
    return (v >= 19968) && (v <= 171941);
  }

  public final boolean incrementToken() throws IOException {
    while (true) {
      if (this.curTermBuffer == null) {
        if (!this.input.incrementToken()) {
          return false;
        }
        if ((!this.nGramNumber) && 
          (isNumeric(this.termAtt.toString()))) {
          return true;
        }
        if ((!this.nGramChinese) && 
          (containsChinese(this.termAtt.toString()))) {
          return true;
        }
        this.curTermBuffer = ((char[])this.termAtt.buffer().clone());

        this.curTermLength = this.termAtt.length();
        this.curGramSize = this.minGram;
        this.tokStart = this.offsetAtt.startOffset();
      }
      if (this.curGramSize <= this.maxGram) {
        if (this.curGramSize >= this.curTermLength) {
          clearAttributes();
          this.offsetAtt.setOffset(this.tokStart + 0, this.tokStart + 
            this.curTermLength);
          this.termAtt.copyBuffer(this.curTermBuffer, 0, 
            this.curTermLength);
          this.curTermBuffer = null;
          return true;
        }
        int start = 0;
        int end = start + this.curGramSize;
        clearAttributes();
        this.offsetAtt.setOffset(this.tokStart + start, this.tokStart + 
          end);
        this.termAtt.copyBuffer(this.curTermBuffer, start, 
          this.curGramSize);
        this.curGramSize += 1;
        return true;
      }

      this.curTermBuffer = null;
    }
  }

  public void reset() throws IOException {
    super.reset();
    this.curTermBuffer = null;
  }

  private boolean isNumeric(String str)
  {
    boolean flag = true;

    if ((str == null) || (str.length() == 0))
      flag = false;
    else {
      for (int i = str.length(); --i >= 0; ) {
        if (!Character.isDigit(str.charAt(i))) {
          flag = false;
        }
      }
    }
    return flag;
  }
}